"""
任务规划
"""
import json
import re
from typing import Any, Dict, List, Optional
from dataclasses import dataclass
from enum import Enum
from langchain.schema import HumanMessage, SystemMessage
import copy
import time

from core.memory import MemoryStore, MemoryType
from core.llm import llm


class PlanType(Enum):
    """智能体可以创建的计划类型。"""
    SEQUENTIAL = "sequential"     # 逐步执行
    PARALLEL = "parallel"         # 并行执行
    CONDITIONAL = "conditional"   # 条件执行
    ITERATIVE = "iterative"       # 迭代执行


@dataclass
class PlanStep:
    """A single step in a plan."""
    id: str
    description: str
    tool: str
    input_template: str
    dependencies: List[str]  # IDs of steps this depends on
    conditions: Optional[Dict[str, Any]] = None  # Conditions for execution
    expected_output: Optional[str] = None
    confidence: float = 0.5
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}


@dataclass
class Plan:
    """一个完整的执行计划。"""
    id: str
    query: str
    goal: str
    plan_type: PlanType
    steps: List[PlanStep]
    estimated_duration: float
    confidence: float
    metadata: Dict[str, Any]
    created_at: float
    
    def get_executable_steps(self, completed_steps: List[str]) -> List[PlanStep]:
        """Get steps that can be executed given completed steps."""
        executable = []
        for step in self.steps:
            if step.id not in completed_steps:
                # Check if all dependencies are completed
                if all(dep in completed_steps for dep in step.dependencies):
                    executable.append(step)
        return executable
    
    def is_complete(self, completed_steps: List[str]) -> bool:
        """检查计划是否完整"""
        return all(step.id in completed_steps for step in self.steps)


class Planner:
    """为复杂查询创建执行计划的高级规划器。"""
    
    def __init__(self, memory_store: MemoryStore = MemoryStore()):
        self.memory_store = memory_store
        self.llm = llm
    
    async def create_plan(self, query: str, available_tools: Optional[List[str]] = None,
                         context: Optional[Dict[str, Any]] = None) -> Plan:
        """为给定的任务制定一个执行计划。"""

        available_tools = available_tools or []
        similar_plans = await self._get_similar_plans(query)
        prompt = self._create_planning_prompt(query, available_tools, similar_plans, context)

        messages = [
            SystemMessage(content=self._get_planning_system_prompt()),
            HumanMessage(content=prompt)
        ]
        response = await self.llm.ainvoke(messages)
        plan_text = response.content

        plan = self._parse_plan(query, plan_text, available_tools)
        await self._store_plan(plan)
        
        return plan
    
    async def refine_plan(self, plan: Plan, execution_results: List[Dict[str, Any]], 
                         current_context: Dict[str, Any]) -> Plan:
        """根据执行结果完善计划。"""
        
        # 分析哪里出了问题，或者有哪些方面可以改进
        refinement_prompt = self._create_refinement_prompt(plan, execution_results, current_context)
        
        messages = [
            SystemMessage(content=self._get_refinement_system_prompt()),
            HumanMessage(content=refinement_prompt)
        ]

        response = await self.llm.ainvoke(messages)
        refinement_text = response.content
        
        # 解析优化内容并更新计划
        refined_plan = self._apply_refinements(plan, refinement_text)
        
        # 存储优化后的计划
        await self._store_plan(refined_plan)
        
        return refined_plan
    
    async def _get_similar_plans(self, query: str, limit: int = 3) -> List[Dict[str, Any]]:
        """Get similar successful plans from memory."""
        similar_memories = await self.memory_store.search_memories(
            query=query,
            memory_type=MemoryType.PLAN_MEMORY,
            limit=limit
        )
        
        successful_plans = []
        for memory in similar_memories:
            if memory.metadata.get("success", False):
                successful_plans.append(memory.content)
        
        return successful_plans
    
    def _create_planning_prompt(self, query: str, available_tools: List[str], 
                               similar_plans: List[Dict[str, Any]], 
                               context: Optional[Dict[str, Any]]) -> str:
        """创建Plan提示词"""
        tools_description = "\n".join([f"- {tool}" for tool in available_tools])
        
        similar_plans_text = ""
        if similar_plans:
            similar_plans_text = "\n\nSimilar successful plans from past:\n"
            for i, plan in enumerate(similar_plans, 1):
                similar_plans_text += f"\nPlan {i}:\n"
                similar_plans_text += f"Query: {plan.get('query', 'N/A')}\n"
                similar_plans_text += f"Steps: {json.dumps(plan.get('steps', []), indent=2)}\n"
        
        context_text = ""
        if context:
            context_text = f"\n\nCurrent context:\n{json.dumps(context, indent=2)}"
        
        return f"""Create a detailed execution plan for the following query:

Query: {query}

Available tools:
{tools_description}
{similar_plans_text}
{context_text}

Please create a plan with the following structure:
1. Goal: Clear statement of what we want to achieve
2. Plan Type: sequential, parallel, conditional, or iterative
3. Steps: List of steps with:
   - ID: unique identifier
   - Description: what this step does
   - Tool: which tool to use
   - Input Template: template for tool input (use {{variable}} for dynamic values)
   - Dependencies: list of step IDs this depends on
   - Expected Output: what we expect to get from this step
   - Confidence: how confident you are this step will work (0.0-1.0)

Format your response as JSON with this structure:
{{
    "goal": "...",
    "plan_type": "sequential|parallel|conditional|iterative",
    "estimated_duration": 30.0,
    "confidence": 0.8,
    "steps": [
        {{
            "id": "step_1",
            "description": "...",
            "tool": "tool_name",
            "input_template": "...",
            "dependencies": [],
            "expected_output": "...",
            "confidence": 0.9
        }}
    ]
}}"""
    
    def _create_refinement_prompt(self, plan: Plan, execution_results: List[Dict[str, Any]], 
                                 current_context: Dict[str, Any]) -> str:
        """Create prompt for plan refinement."""
        
        results_summary = []
        for result in execution_results:
            step_id = result.get("step_id", "unknown")
            success = result.get("success", False)
            error = result.get("error", "")
            results_summary.append(f"Step {step_id}: {'SUCCESS' if success else 'FAILED'} - {error}")
        
        return f"""Analyze the execution results and refine the plan:

Original Plan:
Goal: {plan.goal}
Plan Type: {plan.plan_type.value}
Steps: {json.dumps([step.__dict__ for step in plan.steps], indent=2)}

Execution Results:
{chr(10).join(results_summary)}

Current Context:
{json.dumps(current_context, indent=2)}

Please provide refinements in JSON format:
{{
    "add_steps": [
        {{
            "id": "new_step_1",
            "description": "...",
            "tool": "tool_name",
            "input_template": "...",
            "dependencies": ["step_1"],
            "expected_output": "...",
            "confidence": 0.8
        }}
    ],
    "modify_steps": [
        {{
            "id": "existing_step_id",
            "changes": {{
                "input_template": "new template",
                "dependencies": ["new_deps"]
            }}
        }}
    ],
    "remove_steps": ["step_id_to_remove"],
    "change_plan_type": "sequential|parallel|conditional|iterative"
}}"""
    
    def _parse_plan(self, query: str, plan_text: str, available_tools: List[str]) -> Plan:
        """将大语言模型（LLM）的响应解析为一个计划对象。"""
        try:
            # 从响应中提取JSON
            json_match = re.search(r'\{.*\}', plan_text, re.DOTALL)
            if not json_match:
                raise ValueError("No JSON found in plan response")
            
            plan_data = json.loads(json_match.group())
            
            # 创建计划步骤
            steps = []
            for step_data in plan_data.get("steps", []):
                step = PlanStep(
                    id=step_data.get("id", f"step_{len(steps) + 1}"),
                    description=step_data.get("description", ""),
                    tool=step_data.get("tool", ""),
                    input_template=step_data.get("input_template", ""),
                    dependencies=step_data.get("dependencies", []),
                    expected_output=step_data.get("expected_output"),
                    confidence=step_data.get("confidence", 0.5)
                )
                steps.append(step)

            plan = Plan(
                id=f"plan_{int(time.time())}",
                query=query,
                goal=plan_data.get("goal", ""),
                plan_type=PlanType(plan_data.get("plan_type", "sequential")),
                steps=steps,
                estimated_duration=plan_data.get("estimated_duration", 60.0),
                confidence=plan_data.get("confidence", 0.5),
                metadata={},
                created_at=time.time()
            )
            return plan
            
        except Exception as e:
            # 如果发生异常，则创建一个简单得的序列执行计划。
            return self._create_fallback_plan(query, available_tools)
    
    def _create_fallback_plan(self, query: str, available_tools: List[str]) -> Plan:
        """当解析失败时，制定一个简单的备用计划。"""
        import time

        # 简单启发式规则：如果查询提到计算，使用计算器
        # 如果提到搜索/查找，使用搜索工具，等等。
        steps = []
        
        if any(word in query.lower() for word in ["calculate", "compute", "math"]):
            if "calculator" in available_tools:
                steps.append(PlanStep(
                    id="calc_step",
                    description="Perform calculation",
                    tool="calculator",
                    input_template=query,
                    dependencies=[],
                    confidence=0.7
                ))
        
        if any(word in query.lower() for word in ["search", "find", "look up", "what is"]):
            if "wikipedia" in available_tools:
                steps.append(PlanStep(
                    id="wiki_step",
                    description="Search for information",
                    tool="wikipedia",
                    input_template=query,
                    dependencies=[],
                    confidence=0.6
                ))
        
        # 如果未确定具体步骤，请使用第一个可用的工具
        if not steps and available_tools:
            steps.append(PlanStep(
                id="general_step",
                description="Process query",
                tool=available_tools[0],
                input_template=query,
                dependencies=[],
                confidence=0.4
            ))
        
        return Plan(
            id=f"fallback_plan_{int(time.time())}",
            query=query,
            goal="Process the user query",
            plan_type=PlanType.SEQUENTIAL,
            steps=steps,
            estimated_duration=30.0,
            confidence=0.4,
            metadata={"fallback": True},
            created_at=time.time()
        )
    
    def _apply_refinements(self, plan: Plan, refinement_text: str) -> Plan:
        """对计划进行优化调整。"""
        try:
            json_match = re.search(r'\{.*\}', refinement_text, re.DOTALL)
            if not json_match:
                return plan
            
            refinements = json.loads(json_match.group())

            refined_plan = copy.deepcopy(plan)
            refined_plan.id = f"refined_{plan.id}_{int(time.time())}"

            # 应用优化调整， 添加新步骤
            for step_data in refinements.get("add_steps", []):
                new_step = PlanStep(
                    id=step_data.get("id"),
                    description=step_data.get("description", ""),
                    tool=step_data.get("tool", ""),
                    input_template=step_data.get("input_template", ""),
                    dependencies=step_data.get("dependencies", []),
                    expected_output=step_data.get("expected_output"),
                    confidence=step_data.get("confidence", 0.5)
                )
                refined_plan.steps.append(new_step)
            
            # 修正步骤
            for modification in refinements.get("modify_steps", []):
                step_id = modification.get("id")
                changes = modification.get("changes", {})
                
                for step in refined_plan.steps:
                    if step.id == step_id:
                        for key, value in changes.items():
                            if hasattr(step, key):
                                setattr(step, key, value)
            
            # 移除步骤
            for step_id in refinements.get("remove_steps", []):
                refined_plan.steps = [s for s in refined_plan.steps if s.id != step_id]
            
            # 如果指定了计划类型，则更改计划类型
            if "change_plan_type" in refinements:
                refined_plan.plan_type = PlanType(refinements["change_plan_type"])
            
            return refined_plan

        except Exception:
            return plan
    
    async def _store_plan(self, plan: Plan):
        """在内存中保存Plan"""
        await self.memory_store.remember(
            content={
                "id": plan.id,
                "query": plan.query,
                "goal": plan.goal,
                "plan_type": plan.plan_type.value,
                "steps": [step.__dict__ for step in plan.steps],
                "confidence": plan.confidence
            },
            memory_type=MemoryType.PLAN_MEMORY,
            importance=plan.confidence,
            metadata={
                "plan_id": plan.id,
                "estimated_duration": plan.estimated_duration,
                "step_count": len(plan.steps)
            }
        )
    
    def _get_planning_system_prompt(self) -> str:
        """获取用于规划的系统提示。"""
        return """You are an expert AI planner. Your job is to create detailed, executable plans for complex queries.

Key principles:
1. Break down complex tasks into manageable steps
2. Identify dependencies between steps
3. Choose appropriate tools for each step
4. Provide clear input templates with variable placeholders
5. Estimate confidence levels realistically
6. Consider parallel execution where possible

Plan types:
- Sequential: Steps must be done in order
- Parallel: Some steps can be done simultaneously  
- Conditional: Steps depend on results of previous steps
- Iterative: Steps may need to be repeated

Always respond with valid JSON in the specified format."""
    
    def _get_refinement_system_prompt(self) -> str:
        """获取用于计划优化的系统提示。"""
        return """You are an expert at analyzing execution results and refining plans.

Analyze what went wrong and suggest improvements:
1. Add steps to handle missing functionality
2. Modify steps that failed or produced poor results
3. Remove unnecessary or problematic steps
4. Change plan type if a different approach would work better

Focus on practical improvements that address the specific failures observed.
Always respond with valid JSON in the specified format."""


if __name__ == "__main__":
    import asyncio
    planer = Planner()
    # plan = asyncio.run(planer.create_plan(query='怎么泡制一壶茉莉花茶？'))
    plan = asyncio.run(planer.create_plan(query='计算(2 + 6) * 19 + 5 * 7'))
    print(plan)